from csagent.core.node.base_node import BaseNode
from csagent.core.context import AgentState

class SubAgentNode(BaseNode):
    def execute(self, state: AgentState, stream_writer=None) -> AgentState:
        final_res = state
        for (key, value) in self.graph.stream(state, stream_mode=["custom", "values"]):
            if key == 'values':
                final_res = value
            else:
                stream_writer(value)
        return final_res
    